SageMaker上でDeep Graph LibraryとPyTorchを使用してナレッジグラフ埋め込みの学習を実行する

DA事業本部の貞松です。今回はSageMakerです。
よく考えたら初めてSageMakerの記事を書いているかもしれません。

本記事は『機械学習 on AWS Advent Calendar 2019』9日目のエントリーです。
クラスメソッド 機械学習 on AWS Advent Calendar 2019 - Qiita
クラスメソッド 機械学習 on AWS Advent Calendar 2019 ｜ シリーズ ｜ Developers.IO

先週開催されたAWS re:Invent 2019でSageMakerのアップデートが多数発表されましたが、その中の1つにDeep Graph Library(DGL)がSageMaker上で使用できるようになった、というものがありました。

Amazon Web Services ブログ - Amazon SageMakerですぐに利用可能: Deep Graph Library

本記事では、毎度おなじみ amazon-sagemaker-examples に用意されているサンプルに沿って、DGLの機能や使用する際の手順について解説していきます。

Deep Graph Library(DGL)とは

Deep Graph Library(DGL)とは、Graph Neural Network(GNN)の高速な開発・学習・評価を補助してくれるPythonのオープンソースライブラリです。

DeepGraphLibrary

GNNは、SNS上の人間関係や購買における消費者と商品、化学構造における原子の結合など、ありとあらゆる関係性を示すグラフ構造を持つデータセットに対して予測モデルを作成するために使用されます。

DGLは、PyTorchやMXNetのようなポピュラーなディープラーニングフレームワークの上で動作するようになっています。また、DGL にはグラフのデータセットも含まれており、これを簡単にダウンロードして様々な試行が行えるようになっています。AWSではさらにPyTorchおよびMXNet用のディープラーニングコンテナイメージにDGLを追加しているので、これらのコンテナを使用することでSageMaker上でDGLを簡単に扱えるようになっています。

SageMaker上でDeep Graph Libraryを使ってみる

実行するサンプルの概要

今回使用するサンプルでは、PyTorch用のコンテナイメージを用いて、ナレッジグラフ埋め込み(Knowledge Graph Embedded, KGE)の学習を実行します。

サンプルデータはfb15kというFreebase(Wikipediaなどを基に作成された百科事典データベース)をベースにしたデータを使用します。 コード化された各項目(ページ)とカテゴリの組み合わせでリレーションを表現したデータになっています。

ノートブックインスタンスに対するGitリポジトリ設定

SageMakerのGitリポジトリにamazon-sagemaker-examplesを設定して、ノートブックインスタンス作成時に割り当てておくことで簡単にサンプルを利用することができます。

サンプルノートブックの実行

今回は、リポジトリの amazon-sagemaker-examples/sagemaker-python-sdk/dgl_kge/ に配置されている kge_pytorch.ipynb を使用します。

PyTorch1.3.1用のDockerイメージをECRにプッシュ

このサンプルノートブックを実行するためには、PyTorch用1.3.1Dockerイメージを公式のDeep Learning Containers Imagesから取得して、自前のECRリポジトリにプッシュしておく必要があります。 ※ここの準備に一番時間が掛かります

Deep Learning Containers Images

上記ページの内容に従い、AWS CLIのコマンドを実行してdockerログイン用のコマンドを取得します。

$(aws ecr get-login --no-include-email --region ap-northeast-1 --registry-ids 763104351884)

上記コマンドで出力された以下のようなコマンドをそのままコピペして実行します。

docker login -u AWS -p [とても長い文字列] https://763104351884.dkr.ecr.ap-northeast-1.amazonaws.com

ログインに成功したら、docker pullで必要なdokcerイメージを取得します。

docker pull 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:1.3.1-gpu-py36-cu101-ubuntu16.04

次に自前のECRリポジトリに取得したDockerイメージをプッシュします。

ECRのリポジトリを新規作成し、作成したリポジトリの画面で プッシュコマンドの表示 をクリックします。
表示されたコマンドを参考に順番にコマンドを実行していきます。

まずはAWS CLIのコマンドを実行してECRのdockerログイン用のコマンドを取得します。

$(aws ecr get-login --no-include-email --region ap-northeast-1)

今回は一からDockerイメージを作成するのではなく、既存のイメージを取得しているので、２番目のコマンドは飛ばして３番目のコマンドを実行します。
これにより、取得済みのイメージ名、タグををECR上のリポジトリで設定されているものに合わせます。

docker tag 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:1.3.1-gpu-py36-cu101-ubuntu16.04 xxxxxxxxxxxx.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:latest

最後にdocker pushでイメージをECRリポジトリにプッシュします。

docker push xxxxxxxxxxxx.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-training:latest

SageMakerの設定

ようやくサンプルノートブックを実行できます。
まずは、トレーニングジョブ実行に必要なSageMakerの設定を行います。

import sagemaker
from sagemaker import get_execution_role
from sagemaker.session import Session

# Sagemakerのセッション作成
sess = sagemaker.Session()

# コードとモデルのアーティファクトを保存するS3 bucketの指定 (任意のバケットを指定可能)
bucket = sess.default_bucket()

# カスタムコードの出力先
custom_code_upload_location = 'customcode'

# SagemakerのリソースにアクセスするためのIAMロールの取得
# ノートブック環境からSageMakerのPython SDK経由でロールを取得する
role = get_execution_role()

SageMaker estimator (推定器) の作成

次に実際に学習を実行する部分の設定、作成を行います。

from sagemaker.pytorch import PyTorch

ENTRY_POINT = 'train.py'
CODE_PATH = './'

account = sess.boto_session.client('sts').get_caller_identity()['Account']
region = sess.boto_session.region_name

docker_name = 'pytorch-training' # 前述の作業でECRにプッシュしたコンテナイメージ名を指定
docker_tag = 'latest' # コンテナイメージのタグを指定
image = '{}.dkr.ecr.{}.amazonaws.com/{}:{}'.format(account, region, docker_name, docker_tag)
print(image)

params = {}
params['dataset'] = 'FB15k'  # 
params['model'] = 'DistMult'
params['batch_size'] = 1024
params['neg_sample_size'] = 256
params['hidden_dim'] = 2000
params['gamma'] = 500.0
params['lr'] = 0.1
params['max_step'] = 100000
params['batch_size_eval'] = 16
params['valid'] = True
params['test'] = True
params['neg_adversarial_sampling'] = True

estimator = PyTorch(entry_point=ENTRY_POINT,
                    source_dir=CODE_PATH,
                    role=role, 
                    train_instance_count=1, 
                    train_instance_type='ml.p3.2xlarge',
                    image_name=image,
                    hyperparameters=params,
                    sagemaker_session=sess)

このコードだけ見ると、実際どこでDGLを呼び出して使用しているのか全くわからないので、エントリーポイントである train.py を起点にリポジトリ内のコードを辿っていったところ、MXNetおよびPyTorchのモデル生成の基点になっている general_model.py でライブラリを呼び出しているようです。コード先頭の import dgl.backend as F がそれに該当します。

https://github.com/awslabs/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/dgl_kge/models/general_models.py

この dgl.backend の実体がどういうものであるかは、GitHubのdglリポジトリを参照することで確認できます。

https://github.com/dmlc/dgl/tree/master/python/dgl/backend

Python用のdglライブラリに含まれるbackendモジュールを利用してGNNのモデルを実装する流れですね。

トレーニングジョブの実行

あとは、単純にfit呼び出すことでトレーニングジョブを実行します。

estimator.fit()

トレーニングジョブが完了したらサンプルノートブックの序盤、SageMakerの設定のところで指定したS3バケットに学習済みモデルのアーティファクトが出力されます。 今回の場合はノートブックインスタンスのデフォルトのS3バケットに出力されます。

あとはこれをSageMakerのエンドポイントにデプロイすれば、ナレッジグラフの推論を実行する学習済みモデルを利用可能になります。

まとめ

SageMaker上でDGLを実行するために、最新のDeep Learning Containers Imagesを取得する必要があった為、その点でかなりの時間を取られてしまいましたが、そこさえ越えてしまえば、ノートブックで実行するコードもかなり短く簡単なものでした。今後はSNSのログや購買情報などのもう少しわかりやすいデータを使って、可視化などを含めた実用についても考えたいと思います。